//
// Created by chizuru on 2025/9/29.
//

#include "ModBusProtocolCodec.h"
#include <algorithm>
#include <array>
std::optional<uint64_t> codecArgToUint64(const ProtocolResult::CodecArg &a) {
    return std::visit([](auto&& x) -> std::optional<uint64_t> {
        using T = std::decay_t<decltype(x)>;
        if constexpr (std::is_same_v<T, uint32_t>) return static_cast<uint64_t>(x);
        else if constexpr (std::is_same_v<T, uint64_t>) return x;
        else if constexpr (std::is_same_v<T, double>) return static_cast<uint64_t>(x);
        else if constexpr (std::is_same_v<T, uint16_t>) return static_cast<uint64_t>(x);
        else if constexpr (std::is_same_v<T, uint8_t>) return static_cast<uint64_t>(x);
        else if constexpr (std::is_same_v<T, std::string>) {
            try { return static_cast<uint64_t>(std::stoull(x)); }
            catch (...) { return std::nullopt; }
        } else return std::nullopt;
    }, a);
}

std::vector<uint8_t>
ModBusProtocolCodec::buildWriteFrame(const CommandId &id, const IProtocolCodec::CodecArgs &args) const {
    // args[0] 从站 u8
    // args[1] 起始地址 u16
    // args[2] 寄存器数量 u16
    // args[3] 数据 u16
    // ...
    // args[n-1] 数据 u16
    if (args.size() < 4) {
        return {};
    }

    auto funcOpt = std::get_if<uint8_t>(&id.v);
    if (!funcOpt)
        return {};
    uint8_t func = *funcOpt;
    auto sOpt = codecArgToUint64(args[0]);
    auto addrOpt = codecArgToUint64(args[1]);
    auto qtyOpt = codecArgToUint64(args[2]);
    if (!sOpt || !addrOpt || !qtyOpt) return {};

    uint8_t slave = static_cast<uint8_t>(*sOpt & 0xFF);
    uint16_t startAddr = static_cast<uint16_t>(*addrOpt & 0xFFFF);
    uint16_t qty = static_cast<uint16_t>(*qtyOpt & 0xFFFF);

    // 单寄存器写 (0x06)
    if (func == 0x06) {
        // 要求 qty == 1 且 args[3] 存在
        if (qty != 1) return {}; // 或者返回空表示参数不合规
        auto valOpt = codecArgToUint64(args[3]);
        if (!valOpt) return {};
        uint16_t val16 = static_cast<uint16_t>((*valOpt) & 0xFFFF);

        std::vector<uint8_t> frame;
        frame.reserve(8);
        frame.push_back(slave);
        frame.push_back(func);
        frame.push_back(static_cast<uint8_t>((startAddr >> 8) & 0xFF));
        frame.push_back(static_cast<uint8_t>(startAddr & 0xFF));
        frame.push_back(static_cast<uint8_t>((val16 >> 8) & 0xFF));
        frame.push_back(static_cast<uint8_t>(val16 & 0xFF));
        uint16_t crc = crc16_modbus(frame.data(), frame.size());
        frame.push_back(static_cast<uint8_t>(crc & 0xFF));
        frame.push_back(static_cast<uint8_t>((crc >> 8) & 0xFF));
        return frame;
    }

    // 写多个寄存器 (0x10)
    if (func == 0x10) {
        // 要求提供 qty 个数据（args[3] .. args[3 + qty -1]）
        size_t providedData = args.size() - 3;
        if (providedData < static_cast<size_t>(qty))
            return {}; // 数据不足
        // 构建帧
        std::vector<uint8_t> frame;
        frame.reserve(9 + qty * 2);
        frame.push_back(slave);
        frame.push_back(func);
        frame.push_back(static_cast<uint8_t>((startAddr >> 8) & 0xFF));
        frame.push_back(static_cast<uint8_t>(startAddr & 0xFF));
        frame.push_back(static_cast<uint8_t>((qty >> 8) & 0xFF));
        frame.push_back(static_cast<uint8_t>(qty & 0xFF));
        uint8_t byteCount = static_cast<uint8_t>(qty * 2);
        frame.push_back(byteCount);
        // 每个寄存器取低 16 位写入（按大端顺序）
        for (uint16_t i = 0; i < qty; ++i) {
            auto vOpt = codecArgToUint64(args[3 + i]);
            if (!vOpt) return {};
            uint16_t v16 = static_cast<uint16_t>((*vOpt) & 0xFFFF);
            frame.push_back(static_cast<uint8_t>((v16 >> 8) & 0xFF));
            frame.push_back(static_cast<uint8_t>(v16 & 0xFF));
        }
        uint16_t crc = crc16_modbus(frame.data(), frame.size());
        frame.push_back(static_cast<uint8_t>(crc & 0xFF));
        frame.push_back(static_cast<uint8_t>((crc >> 8) & 0xFF));
        return frame;
    }

    // 不支持的功能码 -> 返回空
    return {};
}

std::vector<uint8_t>
ModBusProtocolCodec::buildReadFrame(const CommandId &id, const IProtocolCodec::CodecArgs &args) const {
    // args[0] 从站地址
    // args[1] 起始地址
    // args[2] 寄存器个数
    if (args.size() < 3)
        return {};

    auto funcOpt = std::get_if<uint8_t>(&id.v);
    if (!funcOpt)
        return {};
    uint8_t func = *funcOpt;

    // 从 args 提取数值
    auto a0 = codecArgToUint64(args[0]);
    auto a1 = codecArgToUint64(args[1]);
    auto a2 = codecArgToUint64(args[2]);
    if (!a0 || !a1 || !a2)
        return {};

    uint8_t slave = static_cast<uint8_t>(*a0 & 0xFF);
    uint16_t startAddr = static_cast<uint16_t>(*a1 & 0xFFFF);
    uint16_t qty = static_cast<uint16_t>(*a2 & 0xFFFF);

    std::vector<uint8_t> frame;
    frame.reserve(8);
    frame.push_back(slave);
    frame.push_back(func);
    // 起始地址高字节、低字节（Modbus 大端寄存器地址）
    frame.push_back(static_cast<uint8_t>((startAddr >> 8) & 0xFF));
    frame.push_back(static_cast<uint8_t>(startAddr & 0xFF));
    // 寄存器数量高、低
    frame.push_back(static_cast<uint8_t>((qty >> 8) & 0xFF));
    frame.push_back(static_cast<uint8_t>(qty & 0xFF));

    // 计算 CRC（对前面的所有字节）
    uint16_t crc = crc16_modbus(frame.data(), frame.size());
    // Modbus 低字节先传
    frame.push_back(static_cast<uint8_t>(crc & 0xFF));
    frame.push_back(static_cast<uint8_t>((crc >> 8) & 0xFF));
    return frame;
}

ProtocolResult
ModBusProtocolCodec::parseWriteResponse(std::span<const uint8_t> resp, const CommandId &commandId, const CodecArgs& args) const {
    // 1 = CRC error
    // 2 = length / byte count mismatch
    // 3 = function code mismatch
    // 4 = parse error / 参数问题
    // 6 = 响应内容与预期不匹配
    // 对于 Modbus 异常（设备返回 0x80|func），返回的 ProtocolResult::Fail(errorCode) 会使用 0x80 | exceptionCode 来携带设备原始异常（便于上层诊断），例如设备返回 0x02 则 errorCode = 0x82。
    if (resp.size() < 5)
        return ProtocolResult::Fail(4);

    // CRC 校验
    size_t n = resp.size();
    uint16_t calc = crc16_modbus(resp.data(), n - 2);
    uint8_t crcLow = resp[n - 2];
    uint8_t crcHigh = resp[n - 1];
    if ((uint8_t)(calc & 0xFF) != crcLow || (uint8_t)((calc >> 8) & 0xFF) != crcHigh) {
        return ProtocolResult::Fail(1);
    }

    auto funcOpt = std::get_if<uint8_t>(&commandId.v);
    if (!funcOpt)
        return ProtocolResult::Fail(4);
    uint8_t expectedFunc = *funcOpt;
    uint8_t respFunc = resp[1];

    // 异常响应 (func | 0x80)
    if (respFunc == (uint8_t)(expectedFunc | 0x80)) {
        // 异常帧通常为 5 字节 (slave, func|0x80, excCode, crcLo, crcHi)
        uint8_t exc = (resp.size() >= 5) ? resp[2] : 0;
        return ProtocolResult::Fail(static_cast<uint8_t>(0x80 | exc));
    }

    if (respFunc != expectedFunc) {
        return ProtocolResult::Fail(3);
    }

    // 对于 0x06 和 0x10 写响应两者标准返回长度均为 8 字节：slave(1) func(1) addr_hi(1) addr_lo(1) qty_or_val_hi(1) qty_or_val_lo(1) crcLo(1) crcHi(1)
    if (resp.size() != 8) {
        return ProtocolResult::Fail(2);
    }

    uint16_t addr = static_cast<uint16_t>((resp[2] << 8) | resp[3]);
    uint16_t qtyOrVal = static_cast<uint16_t>((resp[4] << 8) | resp[5]);

    // 将响应打包为单个 uint64_t 以便上层验证： pack = (addr << 16) | qtyOrVal
    uint64_t pack = (static_cast<uint64_t>(addr) << 16) | static_cast<uint64_t>(qtyOrVal & 0xFFFF);

    // 若上层提供了期望值（args[0]），则进行比较
    if (!args.empty()) {
        auto expectedOpt = codecArgToUint64(args[0]);
        if (!expectedOpt)
            return ProtocolResult::Fail(4);
        uint32_t expectedVal = static_cast<uint32_t>(*expectedOpt);
        if (expectedVal != pack) {
            return ProtocolResult::Fail(6);
        }
    }

    return ProtocolResult::Ok();
}

ProtocolResult ModBusProtocolCodec::parseReadResponse(std::span<const uint8_t> resp, const CommandId &id, const CodecArgs& args) const {
    // 错误码约定（局部实现，仅在本函数内使用）：
    // 1 = CRC error
    // 2 = length / byte count mismatch
    // 3 = function code mismatch or exception
    // 4 = parse error / 参数问题

    if (resp.size() < 5)
        return ProtocolResult::Fail(4); // 最短帧长度：slave(1)+func(1)+bytecnt(1)+crc(2) = 5

    // 校验 CRC（最后两字节为 CRC low, high）
    size_t n = resp.size();
    uint16_t calc = crc16_modbus(resp.data(), n - 2);
    uint8_t crcLow = resp[n - 2];
    uint8_t crcHigh = resp[n - 1];
    if ((uint8_t)(calc & 0xFF) != crcLow || (uint8_t)((calc >> 8) & 0xFF) != crcHigh) {
        return ProtocolResult::Fail(1);
    }

    // 功能码匹配检查
    auto funcOpt = std::get_if<uint8_t>(&id.v);
    if (!funcOpt)
        return ProtocolResult::Fail(4);
    uint8_t expectedFunc = *funcOpt;
    uint8_t respFunc = resp[1];

    // 异常响应：功能码高位置 1（0x80）表示异常，应返回异常码（resp[2]）
    if (respFunc == (uint8_t)(expectedFunc | 0x80)) {
        uint8_t exc = (resp.size() >= 5) ? resp[2] : 0;
        return ProtocolResult::Fail( (uint8_t)(0x80 | exc) ); // 将异常码编码到 errorCode 中（debug）
    }
    if (respFunc != expectedFunc) {
        return ProtocolResult::Fail(3);
    }

    // 字节计数
    uint8_t byteCount = resp[2];
    // 帧长度应为 1+1+1 + byteCount + 2 = 5 + byteCount
    if (resp.size() != static_cast<size_t>(5 + byteCount)) {
        return ProtocolResult::Fail(2);
    }

    const uint8_t* data = resp.data() + 3; // 数据起始
    size_t dataRem = byteCount;
    std::vector<ProtocolResult::CodecArg> values;
    values.reserve(8);

    // 如果 args.empty(), 默认每 2 字节为一个值（uint16 -> 以 uint32 存放）
    if (args.empty()) {
        if (dataRem % 2 != 0) return ProtocolResult::Fail(4);
        size_t count = dataRem / 2;
        for (size_t i = 0; i < count; ++i) {
            uint32_t v = (static_cast<uint32_t>(data[0]) << 8) | static_cast<uint32_t>(data[1]);
            values.emplace_back(static_cast<uint32_t>(v));
            data += 2;
        }
        return ProtocolResult::Ok(std::move(values), std::vector<uint8_t>(resp.begin(), resp.end()));
    }

    // 否则按 args 指定的寄存器个数解析（args[i] 表示第 i 个值占用的寄存器数）
    for (size_t idx = 0; idx < args.size(); ++idx) {
        auto maybeN = codecArgToUint64(args[idx]);
        if (!maybeN)
            return ProtocolResult::Fail(4);
        uint64_t regCount = *maybeN;
        if (regCount == 0)
            return ProtocolResult::Fail(4);

        size_t byteNeeded = static_cast<size_t>(regCount * 2ULL);
        if (dataRem < byteNeeded)
            return ProtocolResult::Fail(2);

        // 将这些字节按 Modbus 大端顺序合成一个 uint32_t（若超过 4 字节则截断或返回错误；按需求这里限制到最多 2 寄存器 -> 4 字节）
        if (byteNeeded > 4) {
            // 超过 4 字节（> 32 bit）目前不支持，返回错误；如果需要支持 uint64 请扩展 CodecArg 和行为
            return ProtocolResult::Fail(4);
        }
        uint32_t val = 0;
        for (size_t b = 0; b < byteNeeded; ++b) {
            val = static_cast<uint32_t>((val << 8) | static_cast<uint32_t>(data[b]));
        }
        // 统一返回 uint32_t（题目要求统一为 uint32_t）
        values.emplace_back(static_cast<uint32_t>(val));
        data += byteNeeded;
        dataRem -= byteNeeded;
    }

    // 如果解析完还有剩余字节，说明 args 与响应不匹配
    if (dataRem != 0) {
        return ProtocolResult::Fail(2);
    }

    return ProtocolResult::Ok(std::move(values), std::vector<uint8_t>(resp.begin(), resp.end()));
}


const std::array<uint16_t, 256> &ModBusProtocolCodec::get_crc16_table() {
    static std::array<uint16_t, 256> table = []() {
        std::array<uint16_t, 256> t{};
        for (int i = 0; i < 256; ++i) {
            uint16_t crc = static_cast<uint16_t>(i);
            for (int j = 0; j < 8; ++j) {
                if (crc & 1) crc = static_cast<uint16_t>((crc >> 1) ^ 0xA001);
                else crc = static_cast<uint16_t>(crc >> 1);
            }
            t[i] = crc;
        }
        return t;
    }();
    return table;
}

uint16_t ModBusProtocolCodec::crc16_modbus(const uint8_t *data, size_t length) {
    const auto &table = get_crc16_table();
    uint16_t crc = 0xFFFF;
    for (size_t i = 0; i < length; ++i) {
        uint8_t idx = static_cast<uint8_t>(crc ^ data[i]);
        crc = static_cast<uint16_t>((crc >> 8) ^ table[idx]);
    }
    return crc;
}

